# -*- coding: utf-8 -*-
# 作者：秋禊
import numpy as np
from PIL import Image
import os
import cv2

xml_path = 'haarcascade_frontalface_default.xml'
face_detector = cv2.CascadeClassifier(xml_path)
recognizer = cv2.face.LBPHFaceRecognizer_create()
path = 'Facedata'

def main01():
    cap = cv2.VideoCapture(0)# 调用笔记本内置摄像头，所以参数为0，如果有其他的摄像头可以调整参数为1，2
    face_id = input('\n 输入用户ID(必须为整数):')
    print('\n 请面向摄像头,正在采集面部信息，请稍后...')
    cai_ji(cap, face_detector, face_id)
    jie_shu(cap)

    faces, ids = get_images_and_labels()
    recognizer.train(faces, np.array(ids))
    recognizer.write(r'face_trainer\trainer.yml')
    print("{0} 个面部。".format(len(np.unique(ids))))
def cai_ji(cap, face_detector, face_id):
    count = 0
    while True:
        sucess, img = cap.read()    # 从摄像头读取图片
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)# 转为灰度图片
        faces = face_detector.detectMultiScale(gray, 1.3, 5) # 检测人脸
        for (x, y, w, h) in faces:
            cv2.rectangle(img, (x,y), (x+w,y+w), (255, 0, 0)) # brief绘制一个矩形
            count += 1
            cv2.imwrite("Facedata/User." + str(face_id) + '.' +
                        str(count) + '.jpg', gray[y: y + h, x: x + w])# 保存图像
            cv2.imshow("lu ru", img)# 初始化一个窗口
        k = cv2.waitKey(1)# 保持画面的持续。
        if k == 27:   # 通过esc键退出摄像
            break
        elif count >= 500:  # 得到500个样本后退出摄像
            break

def get_images_and_labels():
    image_paths = [os.path.join(path, f) for f in os.listdir(path)]  # join函数的作用？# 所有图片路径
    faceSamples = []
    ids = []
    for image_path in image_paths:
        pil_img = Image.open(image_path).convert('L')# 转换为灰度
        # PIL_img = cv2.cvtColor(imagePath, cv2.COLOR_BGR2GRAY)
        img_numpy = np.array(pil_img, 'uint8') # 把灰度图片变为矩阵
        id = int(os.path.split(image_path)[-1].split(".")[1]) # 把图片路径，以 。为界分割，第一项为user的名字
        faces = face_detector.detectMultiScale(img_numpy) # 把检测到的对象图像大小的值，作为矩阵返回
        for (x, y, w, h) in faces:
            faceSamples.append(img_numpy[y: y+h, x: x+w])
            ids.append(id)
    return faceSamples, ids

def shi_bie():
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    recognizer.read('face_trainer/trainer.yml')
    font = cv2.FONT_HERSHEY_SIMPLEX  # 字体

    idnum = 0
    names = ['one', 'two']
    # 调用笔记本内置摄像头，所以参数为0，如果有其他的摄像头可以调整参数为1，2
    cam = cv2.VideoCapture(0)
    # 设置大小
    minW = 0.1 * cam.get(3)
    minH = 0.1 * cam.get(4)

    while True:
        # 从摄像头读取图片
        ret, img = cam.read()
        if ret == False:
            break
            jie_shu(cam)
        elif ret == True:
            # 把读取图片，转换为灰度图片
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            # 图片中的每一张人脸大小不同，矩形框也不一样大，把检测到的矩形框大x,y,w,h, 列表形式返回
            # 如果没有检测到面部，返回一个元组
            faces = face_detector.detectMultiScale(
                gray,
                scaleFactor=1.2,  # 指定在每个图像尺度上缩小多少
                minNeighbors=5,  # 指定每个候选矩形应该保留多少个邻居。
                minSize=(int(minW), int(minH))  # 最小对象大小的值。小于该值的对象将被忽略
                # minSize和maxSize用来限制得到的目标区域的范围
            )
            # 绘制矩形
            # 图片的x轴、y轴位置，w表示宽、h表示高
            for (x, y, w, h) in faces:
                cv2.rectangle(img,  # 矩形
                              pt1=(x, y),  # 矩形的顶点
                              pt2=(x + w, y + h),  # pt1的对角点
                              color=(0, 255, 0),  # 矩形颜色或亮度(灰度图像)。
                              thickness=2  # 矩形边框的厚度
                              )
                # 返回侦测到的人脸的id和近似度conf（数字越大和训练数据越不像）
                idnum, confidence = recognizer.predict(gray[y:y + h, x:x + w])

                if confidence < 100:
                    idnum = names[idnum]
                    confidence = "{0}%".format(round(100 - confidence))
                else:
                    idnum = "unknown"
                    confidence = "{0}%".format(round(100 - confidence))
                # 方框上的名字
                cv2.putText(img, str(idnum), (x + 5, y - 5), font, 1, (255, 0, 0), 1)
                # 百分比
                cv2.putText(img, str(confidence), (x + 5, y + h - 5), font, 1, (255, 255, 255), 1)

            # 显示指定窗口中的图像
            cv2.imshow('Face recognition', img)
            # 按 Esc 键退出
            k = cv2.waitKey(10)
            if k == 27:
                break
            jie_shu(cam)

def jie_shu(cap):
    cap.release()# 关闭摄像头，释放内存
    cv2.destroyAllWindows()# 关闭窗口

def main():
    print("按1：录入面部\n按2：识别面部\n按3：退出")
    while True:
        num = int(input("请输入："))
        if num == 1:
            main01()
        elif num == 2:
            shi_bie()
        elif num == 3:
            break

if __name__ == '__main__':
    main()